package com.lambda.plugin;

import com.android.build.api.transform.Context;
import com.android.build.api.transform.DirectoryInput;
import com.android.build.api.transform.Format;
import com.android.build.api.transform.JarInput;
import com.android.build.api.transform.QualifiedContent;
import com.android.build.api.transform.Transform;
import com.android.build.api.transform.TransformException;
import com.android.build.api.transform.TransformInput;
import com.android.build.api.transform.TransformInvocation;
import com.android.build.api.transform.TransformOutputProvider;
import com.android.build.gradle.internal.pipeline.TransformManager;

import org.apache.commons.codec.digest.DigestUtils;
import org.apache.commons.io.FileUtils;
import org.objectweb.asm.AnnotationVisitor;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.commons.AdviceAdapter;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Set;

public class ASMTransform extends Transform {
    private final ASMConfig mConfig;

    ASMTransform(ASMConfig config) {
        mConfig = config;
    }

    @Override
    public String getName() {
        return ASMTransform.class.getSimpleName();
    }

    @Override
    public Set<QualifiedContent.ContentType> getInputTypes() {
        return TransformManager.CONTENT_CLASS;
    }

    @Override
    public Set<? super QualifiedContent.Scope> getScopes() {
        return TransformManager.SCOPE_FULL_PROJECT;
    }

    @Override
    public boolean isIncremental() {
        return false;
    }

    @Override
    public void transform(TransformInvocation transformInvocation) throws TransformException, InterruptedException, IOException {
        _transform(transformInvocation.getContext(), transformInvocation.getInputs(),
                transformInvocation.getReferencedInputs(),
                transformInvocation.getOutputProvider(),
                transformInvocation.isIncremental());
    }

    private void _transform(
            Context context,
            Collection<TransformInput> inputs,
            Collection<TransformInput> referencedInputs,
            TransformOutputProvider outputProvider,
            boolean isIncremental) throws IOException, TransformException, InterruptedException {
        if (!isIncremental) {
            outputProvider.deleteAll();
        }

        for (String pakg : mConfig.pakgs) System.out.println(pakg);

        for (TransformInput transformInput : inputs) {
            for (DirectoryInput directoryInput : transformInput.getDirectoryInputs()) {
                File destDir = outputProvider.getContentLocation(
                        directoryInput.getName(),
                        directoryInput.getContentTypes(),
                        directoryInput.getScopes(),
                        Format.DIRECTORY
                );
                FileUtils.copyDirectory(directoryInput.getFile(), destDir);
                traverse(destDir, directoryInput.getFile());
            }

            for (JarInput jarInput : transformInput.getJarInputs()) {
                String destName = jarInput.getName();
                if (destName.endsWith(".jar")) {
                    destName = destName.substring(0, destName.length() - 4);
                }

                String hexName = DigestUtils.md5Hex(jarInput.getFile().getAbsolutePath());
                File destDir = outputProvider.getContentLocation(
                        destName + "_" + hexName,
                        jarInput.getContentTypes(),
                        jarInput.getScopes(),
                        Format.JAR
                );

                FileUtils.copyFile(jarInput.getFile(), destDir);
            }
        }
    }

    private void traverse(File dest, File file) {
        List<File> buff = new ArrayList<>();
        buff.add(file);

        final String path = file.getAbsolutePath();

        while (!buff.isEmpty()) {
            File tmp = buff.remove(0);

            if (!tmp.isDirectory()) {
                String classPath = tmp.getAbsolutePath().substring(path.length() + 1);
                String className = classPath.replace(File.separatorChar, '.');
                if (packageTest(className)) {
                    File target = new File(dest + "/" + classPath);
                    try (
                            FileInputStream fis = new FileInputStream(tmp);
                            FileOutputStream fos = new FileOutputStream(target)
                    ) {
                        ClassReader cr = new ClassReader(fis);
                        ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_MAXS);
                        cr.accept(new MMClassVisitor(cw), ClassReader.EXPAND_FRAMES);

                        fos.write(cw.toByteArray());
                    } catch (Exception e) {
                        e.printStackTrace(System.out);
                    }
                }
            } else {
                File[] files = tmp.listFiles();
                if (files != null) {
                    buff.addAll(Arrays.asList(files));
                }
            }
        }
    }

    private boolean packageTest(String targetPakg) {
        for (String pakg : mConfig.pakgs) {
            if (targetPakg.startsWith(pakg)) {
                String tt = targetPakg.substring(pakg.length() + 1);
                return !tt.startsWith("R$") &&
                        !tt.startsWith("R.class") &&
                        !tt.startsWith("BuildConfig");
            }
        }
        return false;
    }

    static class MMClassVisitor extends ClassVisitor {

        MMClassVisitor(ClassVisitor visitor) {
            super(Opcodes.ASM5, visitor);
        }

        @Override
        public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
            MethodVisitor visitor = super.visitMethod(access, name, desc, signature, exceptions);
            return new MMMethodVisitor(visitor, access, name, desc);
        }
    }

    static class MMMethodVisitor extends AdviceAdapter {
        private final MethodVisitor mMethodVisitor;
        private AnnVisitor mAnnVisitor;
        private boolean mChecked = false;

        MMMethodVisitor(MethodVisitor mv, int access, String name, String desc) {
            super(Opcodes.ASM5, mv, access, name, desc);
            mMethodVisitor = mv;
        }

        @Override
        public AnnotationVisitor visitAnnotation(String desc, boolean visible) {
            // ## Visit Annotation: Lcom/example/app/Aspect;
            AnnotationVisitor tav = super.visitAnnotation(desc, visible);
            if ("Lcom/iyoyi/library/ASM;".equals(desc)) {
                mChecked = true;
                mAnnVisitor = new AnnVisitor(Opcodes.ASM5, tav);
                return mAnnVisitor;
            } else {
                return tav;
            }
        }

        @Override
        protected void onMethodEnter() {
            if (!mChecked) return;

            mMethodVisitor.visitLdcInsn(mAnnVisitor.getLabel());
            mMethodVisitor.visitLdcInsn(mAnnVisitor.getDescribe());
            mMethodVisitor.visitMethodInsn(INVOKESTATIC,
                    "android/util/Log", "e",
                    "(Ljava/lang/String;Ljava/lang/String;)I",
                    false
            );
        }
    }

    static class AnnVisitor extends AnnotationVisitor {
        private final HashMap<String, Object> keyValue = new HashMap<>();

        AnnVisitor(int api, AnnotationVisitor av) {
            super(api, av);
        }

        @Override
        public void visit(String name, Object value) {
            super.visit(name, value);
            keyValue.put(name, value);
        }

        public String getLabel() {
            return (String) keyValue.get("label");
        }

        public Object getDescribe() {
            return keyValue.get("describe");
        }
    }
}
